from picamera2 import Picamera2
import cv2
from threading import Thread
import uuid
import os
import time

count = 0

def image_collect(picam):
    global count
    while True:
        # 捕获帧并转换颜色
        array = picam.capture_array("main")
        img = cv2.cvtColor(array, cv2.COLOR_RGB2BGR)
        
        file_name = str(uuid.uuid4()) + '.jpg'
        cv2.imwrite(os.path.join('images', file_name), img)
        count += 1
        print("save %d %s" % (count, file_name))
        time.sleep(0.4)  # 控制保存频率

if __name__ == "__main__":
    # 创建保存目录
    os.makedirs("images", exist_ok=True)

    # 初始化 Picamera2
    picam2 = Picamera2()
    picam2.configure(picam2.create_video_configuration(
        main={"size": (640, 480)},
        controls={"FrameRate": 30.0}
    ))
    picam2.start()
    time.sleep(1)

    # FPS 计算变量
    fps_start = time.time()
    frame_count = 0
    fps = 0.0
    display_interval = 3  # 每 3 帧显示一次

    # 创建捕获线程
    m_thread = Thread(target=image_collect, args=([picam2]), daemon=True)

    try:
        while True:
            # 捕获帧
            array = picam2.capture_array("main")
            
            frame_count += 1
            elapsed_time = time.time() - fps_start
            if elapsed_time >= 1.0:
                fps = frame_count / elapsed_time
                frame_count = 0
                fps_start = time.time()

            # 只在特定间隔显示
            if frame_count % display_interval == 0:
                img = cv2.cvtColor(array, cv2.COLOR_RGB2BGR)
                small_img = cv2.resize(img, (320, 240))  # 缩小显示
                cv2.putText(small_img, f"FPS: {fps:.1f}", 
                           (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                           0.7, (0, 255, 0), 1)
                cv2.imshow("video", small_img)
                cv2.waitKey(1)

            key = cv2.waitKey(1) & 0xFF
            if key == ord('c'):
                if not m_thread.is_alive():  # 避免重复启动
                    m_thread.start()
            elif key == ord('q'):
                break

    except KeyboardInterrupt:
        print("\nStopped by user")

    finally:
        picam2.stop()
        cv2.destroyAllWindows()
        print("Camera and windows closed")
